package com.mall.gateway.config;

import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayFlowRule;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayRuleManager;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.GatewayCallbackManager;
import com.alibaba.fastjson.JSON;
import com.mall.common.result.Result;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.PostConstruct;
import java.util.HashSet;
import java.util.Set;

/**
 * Sentinel配置类
 */
@Configuration
public class SentinelConfig {

    /**
     * 初始化限流规则
     */
    @PostConstruct
    public void initGatewayRules() {
        Set<GatewayFlowRule> rules = new HashSet<>();
        
        // 对所有请求进行限流，每秒最多1000个请求
        GatewayFlowRule rule = new GatewayFlowRule()
                .setResource("mall-gateway")
                .setCount(1000)
                .setIntervalSec(1);
        rules.add(rule);
        
        // 对登录接口进行限流，每秒最多100个请求
        GatewayFlowRule loginRule = new GatewayFlowRule()
                .setResource("/auth/login")
                .setCount(100)
                .setIntervalSec(1);
        rules.add(loginRule);
        
        // 对注册接口进行限流，每秒最多10个请求
        GatewayFlowRule registerRule = new GatewayFlowRule()
                .setResource("/auth/register")
                .setCount(10)
                .setIntervalSec(1);
        rules.add(registerRule);
        
        // 对获取验证码接口进行限流，每秒最多100个请求
        GatewayFlowRule captchaRule = new GatewayFlowRule()
                .setResource("/auth/captcha")
                .setCount(100)
                .setIntervalSec(1);
        rules.add(captchaRule);
        
        // 对发送短信验证码接口进行限流，每秒最多10个请求
        GatewayFlowRule smsCodeRule = new GatewayFlowRule()
                .setResource("/auth/sms/code")
                .setCount(10)
                .setIntervalSec(1);
        rules.add(smsCodeRule);
        
        GatewayRuleManager.loadRules(rules);
        
        // 设置限流回调处理
        GatewayCallbackManager.setBlockHandler(new BlockRequestHandler() {
            @Override
            public Mono<ServerResponse> handleRequest(ServerWebExchange exchange, Throwable t) {
                Result<Object> result = Result.failed("请求过于频繁，请稍后再试");
                String resultStr = JSON.toJSONString(result);
                
                return ServerResponse.status(HttpStatus.TOO_MANY_REQUESTS)
                        .contentType(MediaType.APPLICATION_JSON)
                        .bodyValue(resultStr);
            }
        });
    }
}